package com.leetcode.partition7;

import java.util.HashMap;
import java.util.Map;

/**
 * @author `RKC`
 * @date 2021/11/14 9:11
 */
public class LC677键值映射 {

    public static void main(String[] args) {
        MapSum mapSum = new MapSum();
        mapSum.insert("abc", 1);
        mapSum.insert("abd", 1);
        mapSum.insert("abe", 1);
        mapSum.insert("abef", 1);
        mapSum.insert("a", 2);
        System.out.println(mapSum.sum("ab"));
    }

    private static class MapSum {
        private final TrieNode root = new TrieNode(' ', 0);

        public MapSum() { }

        public void insert(String key, int val) {
            TrieNode cur = root;
            for (int i = 0; i < key.length(); i++) {
                if (!cur.childNodes.containsKey(key.charAt(i))) {
                    TrieNode n = new TrieNode(key.charAt(i), val);
                    cur.childNodes.put(key.charAt(i), n);
                }
                cur = cur.childNodes.get(key.charAt(i));
            }
            cur.isKey = true;
            cur.val = val;
        }

        public int sum(String prefix) {
            TrieNode cur = root;
            for (int i = 0; i < prefix.length(); i++) {
                if (!cur.childNodes.containsKey(prefix.charAt(i))) return 0;
                cur = cur.childNodes.get(prefix.charAt(i));
            }
            return sum(cur);
        }

        private int sum(TrieNode root) {
            if (root == null) return 0;
            int res = 0;
            if (root.isKey) res += root.val;
            for (Map.Entry<Character, TrieNode> entry : root.childNodes.entrySet()) {
                res += sum(entry.getValue());
            }
            return res;
        }
    }

    private static class TrieNode {
        private char key;
        private Map<Character, TrieNode> childNodes;
        private boolean isKey;
        private int val;

        public TrieNode(char key, int val) {
            this.key = key;
            this.val = val;
            childNodes = new HashMap<>();
            isKey = false;
        }
    }
}
